import os

import SimpleITK as sitk
from PIL import Image
import scipy.ndimage as ndimage
import numpy as np

'''

BTCV的标签
背景：0
肝（liver）：6
右肾（right kidney）：2
左肾（left kidney）：3
脾（spleen）：1
'''
def getLabel(label_array):
    '''
    1：肝脏
    2：肿瘤
    '''
    # label_array[label_array==1] = 0
    # label_array[label_array==2]=1
    return label_array
def getMaxAndMinFromNii(image_array,percentage=0.98):

    '''
    阈值截断，大于percentage的值设置为percentage的值设置为98的值
    '''
    n_array=image_array.flatten()
    n_array=np.sort(n_array)#升序
    p_max=n_array[int(len(n_array)*percentage):].min()

    image_array[image_array>p_max]=p_max
    return image_array
def getSlice(label_array,n_class=2):
    '''
    获取三个方向上包含目标的切片的起始位置
    label_array:标注array
    n_class：类别数,不包含背景
    45 74     142 369     70 379
    45 74     149 398     66 390
    295 476     157 407     49 361
    '''

    spand_slice_z=0
    spand_slice_lx = 3
    spand_slice_rx = 3
    spand_slice_ly = 3
    spand_slice_ry = 3
    contain_class=None
    for i in range(1,n_class+1):
        if i==1:
            contain_class = np.any(label_array == i, axis=(1, 2))
        else:
            contain_class=np.logical_or(np.any(label_array == i, axis=(1, 2)), contain_class)
    indexs_z = list(np.where(contain_class))[0]
    start_slice_z = max(0, indexs_z[0] - spand_slice_z)
    end_slice_z = min(label_array.shape[0], indexs_z[-1] + spand_slice_z) + 1

    for i in range(1, n_class + 1):
        if i == 1:
            contain_class = np.any(label_array == i, axis=(0, 2))
        else:
            contain_class = np.logical_or(np.any(label_array == i, axis=(0, 2)), contain_class)
    indexs_x = list(np.where(contain_class))[0]
    start_slice_x = max(0, indexs_x[0] - spand_slice_lx)
    end_slice_x = min(label_array.shape[1], indexs_x[-1] + spand_slice_rx) + 1

    for i in range(1, n_class + 1):
        if i == 1:
            contain_class = np.any(label_array == i, axis=(0, 1))
        else:
            contain_class = np.logical_or(np.any(label_array == i, axis=(0, 1)), contain_class)
    indexs_y = list(np.where(contain_class))[0]
    start_slice_y = max(0, indexs_y[0] - spand_slice_ly)
    end_slice_y = min(label_array.shape[2], indexs_y[-1] + spand_slice_ry) + 1
    return start_slice_z,end_slice_z,start_slice_x,end_slice_x,start_slice_y,end_slice_y
def processAndSaveATLAS(data_dir=r"/home/liukai/AllData/ATLAS2023/train",
                   save_dir=r"/home/liukai/AllData/liverTumorForDomain"):
    if not os.path.exists(save_dir):
        os.makedirs(save_dir)
    number = 0
    spacing_z=3#z轴的分辨率统一
    label_dir = os.path.join(data_dir,"labelsTr" )
    image_dir = os.path.join(data_dir,"imagesTr" )
    subject_list=os.listdir(image_dir)
    subject_list.sort()
    for index,nii_name in (enumerate(subject_list)):
        #---读取图像数据----#
        image_path = os.path.join(image_dir, nii_name)
        image_nii = sitk.ReadImage(image_path)
        image_array = sitk.GetArrayFromImage(image_nii)
        label_path = os.path.join(label_dir, nii_name.replace('im', 'lb'))
        label_nii = sitk.ReadImage(label_path)
        label_array = sitk.GetArrayFromImage(label_nii)
        # print("提取目标区域前的shape",label_array.shape, image_array.shape)
        # -------------------提取image和label的肝脏和肿瘤区域-----------#
        start_slice_z, end_slice_z, start_slice_x, end_slice_x, start_slice_y, end_slice_y = getSlice(label_array)
        image_array = image_array[start_slice_z:end_slice_z, start_slice_x:end_slice_x, start_slice_y:end_slice_y]
        label_array = label_array[start_slice_z:end_slice_z, start_slice_x:end_slice_x, start_slice_y:end_slice_y]
        # print("提取目标区域后的shape",label_array.shape, image_array.shape)
        # -------------------图像数据阈值截断---------------------#
        image_array=getMaxAndMinFromNii(image_array,percentage=0.98)
        # -------------------处理图像数据的spacing---------------------#
        last_shape=np.array([len(image_array),256,256])
        image_dimension_adjustment = 1 / (image_array.shape / last_shape)
        image_dimension_adjustment[0]=image_nii.GetSpacing()[-1] / spacing_z
        image_array = ndimage.zoom(image_array, image_dimension_adjustment, order=3)  # 双线性插值
        #---z-score   数值放缩到[-1,1]之间，2 * ((array - m) / (l - (m))) - 1--#
        image_array = (image_array - image_array.mean()) / image_array.std()
        image_array=2*((image_array-image_array.min())/(image_array.max()-image_array.min()))-1

        # -------------------处理标注数据---------------------#
        label_array=getLabel(label_array)
        # print(set(label_array.flatten()))
        label_dimension_adjustment = 1 / (label_array.shape / last_shape)
        label_dimension_adjustment[0] = label_nii.GetSpacing()[-1] / spacing_z
        label_array = ndimage.zoom(label_array, label_dimension_adjustment, order=0)  # 只会出现原来出现过的数值
        print("调整spacing后的shape:",index,image_array.shape,label_array.shape)
        #--------拼接image和label的每个2D切片，[image, mask] 。并保存-----#
        for i in range(len(image_array)):
            if 1 in label_array[i]:
                number += 1
                # 拼接并保存
                # combine_array = np.array([image_array[i], label_array[i]])
                last_dir=os.path.join(save_dir,nii_name.split('.')[0])
                if not os.path.exists(last_dir):
                    os.makedirs(last_dir)
                save_name = os.path.join(last_dir, str(i+1) + '.npz')
                np.savez(save_name, arr_0=image_array[i],arr_1=label_array[i])#为了跟代码的dateset对应，参数名和dataset统一arr_0，arr_1
    print(number)
def processAndSaveLITS(data_dir=r"/home/liukai/AllData/ATLAS2023/train",
                   save_dir=r"/home/liukai/AllData/liverTumorForDomain"):
    if not os.path.exists(save_dir):
        os.makedirs(save_dir)
    number = 0
    spacing_z=2#z轴的分辨率统一
    label_dir = os.path.join(data_dir,"segmentation" )
    image_dir = os.path.join(data_dir,"volume" )
    subject_list=os.listdir(image_dir)
    subject_list.sort()
    #阈值截断，看别的论文LITS数据集一般设为[-200,250]
    up=250
    down=-200
    for index,nii_name in enumerate(subject_list):
        #---读取图像数据----#
        image_path = os.path.join(image_dir, nii_name)
        image_nii = sitk.ReadImage(image_path)
        image_array = sitk.GetArrayFromImage(image_nii)
        label_path = os.path.join(label_dir, nii_name.replace('volume', 'segmentation'))
        label_nii = sitk.ReadImage(label_path)
        label_array = sitk.GetArrayFromImage(label_nii)
        # print("提取目标区域前的shape",label_array.shape, image_array.shape)
        # -------------------提取image和label的肝脏和肿瘤区域-----------#
        start_slice_z, end_slice_z, start_slice_x, end_slice_x, start_slice_y, end_slice_y = getSlice(label_array)
        image_array = image_array[start_slice_z:end_slice_z, start_slice_x:end_slice_x, start_slice_y:end_slice_y]
        label_array = label_array[start_slice_z:end_slice_z, start_slice_x:end_slice_x, start_slice_y:end_slice_y]
        # print("提取目标区域后的shape",label_array.shape, image_array.shape)
        # -------------------处理图像数据---------------------#
        last_shape=np.array([len(image_array),256,256])
        image_dimension_adjustment = 1 / (image_array.shape / last_shape)
        image_dimension_adjustment[0]=image_nii.GetSpacing()[-1] / spacing_z
        image_array = ndimage.zoom(image_array, image_dimension_adjustment, order=3)  # 双线性插值
        image_array[image_array>up]=up
        image_array[image_array<down]=down
        image_array = (image_array - image_array.mean()) / image_array.std()
        print(image_array.max(), image_array.min())
        image_array = 2 * ((image_array - image_array.min()) / (image_array.max() - image_array.min())) - 1

        # -------------------处理标注数据---------------------#
        label_array=getLabel(label_array)
        # print(set(label_array.flatten()))
        label_dimension_adjustment = 1 / (label_array.shape / last_shape)
        label_dimension_adjustment[0] = image_nii.GetSpacing()[-1] / spacing_z
        label_array = ndimage.zoom(label_array, label_dimension_adjustment, order=0)  # 只会出现原来出现过的数值
        print("调整spacing后的shape:",index,image_array.shape,label_array.shape)
        # print()
        print(index)
        #--------拼接image和label的每个2D切片，[image, mask] 。并保存-----#
        for i in range(len(image_array)):
            if 1 in label_array[i]:
                number += 1
                # 拼接并保存
                # combine_array = np.array([image_array[i], label_array[i]])
                last_dir=os.path.join(save_dir,nii_name.split('.')[0])
                if not os.path.exists(last_dir):
                    os.makedirs(last_dir)
                save_name = os.path.join(last_dir, str(i+1) + '.npz')
                np.savez(save_name, arr_0=image_array[i],arr_1=label_array[i])#为了跟代码的dateset对应，参数名和dataset统一arr_0，arr_1
    print(number)
if __name__ == '__main__':
    #gpu2
    processAndSaveATLAS(data_dir=r"/home/liukai/AllData/ATLAS2023/train",
                   save_dir=r"/home/liukai/AllData/liverTumorForDomain/ATLAS2023AfterProcess")
    processAndSaveLITS(data_dir=r"/home/liukai/AllData/LITS",
                        save_dir=r"/home/liukai/AllData/liverTumorForDomain/LITSAfterProcess")
    # 本机
    # processAndSaveATLAS(data_dir=r"D:\AllData\atlas-train-dataset-1.0.1\atlas-train-dataset-1.0.1\train",
    #                     save_dir=r"D:\AllData\atlas-train-dataset-1.0.1\atlas-train-dataset-1.0.1\cutslice")
    # processAndSaveLITS(data_dir=r"D:\AllData\LITS",
    #                     save_dir=r"D:\AllData\LITS\cutSlice")
